# -*- coding: utf-8 -*-
"""
Created on Sat Oct  3 21:25:37 2020

@author: jyh
"""
import tensorflow as tf
import numpy as np
from sklearn.preprocessing import OneHotEncoder
import xlwt
import re
import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'


# =============================================================================
# 此实验旨在将 ACARS FTP_SSH AIS TLS几个协议数据导进来，组合实验进行探索
# =============================================================================

def data_under_test(path, n_acars_up_encr=0,n_acars_down_encr=0,n_tls=0,n_ssh=0,n_ais_en=0,n_acars_down_pain=0,n_acars_up_pain=0,n_ais_pain=0,n_http=0,n_dns=0,L_mesg=40):
    #L_mesg 最长数据报文的长度,所有数据最长为40字节
    '''''''''数据处理'''''''''
    
    """----读取加密数据----"""
    '''读取AIS加密数据'''
    ais_en = np.load(os.path.join(path, 'dataset', 'whole_frame', 'AIS_aes_dec_38head_128load_2tail.npy')) # 150460
    '''读取TLS数据'''
    tls = np.load(os.path.join(path, 'dataset', 'whole_frame', 'tls_maccdc_dec_5head_35payload.npy'))  # 136389
    '''读取FTP_SSH数据'''
    tcp_ssh = np.load(os.path.join(path, 'dataset', 'whole_frame', 'tcp_ssh_maccdc_dec_16head_24load.npy'))    # 9479
    '''读取ACARS加密数据'''
    acars_down_en = np.load(os.path.join(path, 'dataset', 'whole_frame', 'acars_down_aes_dec_24head_16load.npy'))  # 23919
    acars_up_en = np.load(os.path.join(path, 'dataset', 'whole_frame', 'acars_up_aes_dec_13head_16load.npy'))  # 25741
    
    """----读取非加密数据----"""
    '''读取ACARS非加密数据'''
    acars_down_un= np.load(os.path.join(path, 'dataset', 'whole_frame', 'racars_down_pain_dec_withhead(40B).npy')) # 31378
    acars_up_un = np.load(os.path.join(path, 'dataset', 'whole_frame', 'racars_up_pain_dec_withhead(40B).npy'))    # 45055
    '''读取AIS非加密数据'''
    ais_un = np.load(os.path.join(path, 'dataset', 'whole_frame', 'AIS_pain_dec_xiaoxi1_with_head.npy'))   # 150460
    '''读取HTTP数据'''
    http = np.load(os.path.join(path, 'dataset', 'whole_frame', 'http_whole_fram_dec_payload(64B).npy'))   # 128249
    http = http[:,0:L_mesg]
    '''读取DNS数据'''
    dns = np.load(os.path.join(path, 'dataset', 'whole_frame', 'dns_maccdc_dec_40B.npy'))  # 7175
    
    #data = np.vstack((up_aes,down_aes[:,0:29])) #上下文混合后的数据进行截短对齐
    ais_en_pad = np.hstack((ais_en, 100*np.ones([len(ais_en),L_mesg-21],float)))
    ais_un_pad = np.hstack((ais_un, 100*np.ones([len(ais_un),L_mesg-21],float)))
    acars_up_en_pad = np.hstack((acars_up_en,100*np.ones([len(acars_up_en),L_mesg-29],float))) #上下文混合后的数据进行补0对齐
    
      
    d_dict1 = {1:acars_up_en_pad[0:n_acars_up_encr],2:acars_down_en[0:n_acars_down_encr],3:tls[0:n_tls],4:tcp_ssh[0:n_ssh],5:ais_en_pad[0:n_ais_en],\
              6:acars_down_un[0:n_acars_down_pain],7:acars_up_un[0:n_acars_up_pain],8:ais_un_pad[0:n_ais_pain],9:http[0:n_http],10:dns[0:n_dns]}
    
    f_dict1 = {1:'ACARS_up_en_',2:'ACARS_down_en_',3:'TLS_',4:'TCPSSH_',5:'AIS_en_',6:'ACARS_down_un_',7:'ACARS_up_un_',8:'AIS_un_',9:'HTTP_',10:'DNS_'}
    
    temp = [n_acars_up_encr, n_acars_down_encr, n_tls, n_ssh, n_ais_en, n_acars_down_pain,n_acars_up_pain,n_ais_pain,n_http,n_dns]
    data_choose = []#选择要组合的数据
    for i in range(1,11):
        if temp[i-1]!=0:
            data_choose.append(i)

    file_out = ""
    data_pad = np.array([])
    data_pad = np.reshape(data_pad,[0,L_mesg])
    for item in data_choose:
        file_out = file_out + f_dict1[item]
        data_pad = np.vstack((data_pad,d_dict1[item]))
    np.random.shuffle(data_pad)
    
    f_dict2 = {1:'ACARS_up_encr',2:'ACARS_down_encr',3:'TLS',4:'SSH',5:'AIS_encr',6:'ACARS_down_plain',7:'ACARS_up_plain',8:'AIS_plain',9:'HTTP',10:'DNS'}
    data = {'ACARS_up_encr':acars_up_en[0:n_acars_up_encr],'ACARS_down_encr':acars_down_en[0:n_acars_down_encr],'TLS':tls[0:n_tls],'SSH':tcp_ssh[0:n_ssh],\
            'AIS_encr':ais_en[0:n_ais_en],'ACARS_up_plain':acars_up_un[0:n_acars_up_pain],'ACARS_down_plain':acars_down_un[0:n_acars_down_pain],\
                'AIS_plain':ais_un[0:n_ais_pain],'HTTP':http[0:n_http],'DNS':dns[0:n_dns]}
   
    temp1 = {'ACARS_up_encr':[13,28],'ACARS_down_encr':[24,39],'TLS':[5,39],'SSH':[16,39],'AIS_encr':[5,20],
             'ACARS_up_plain':[-1,-1],'ACARS_down_plain':[-1,-1],'AIS_plain':[-1,-1],'HTTP':[-1,-1],'DNS':[-1,-1]}  #加密段起止位置，python索引从0开始，非加密数据设为[-1,-1]
    
    sel_pro_name = {}
    sel_data = {}
    sel_label = {}
    for item in data_choose:
        temp = f_dict2[item]
        sel_pro_name[1]=temp
        sel_data[temp]=data[temp]
        sel_label[temp]=temp1[temp]
    
    label = {}
    for i in sel_data.keys():
        temp2 = []
        for j in range(len(sel_data[i])):
            temp2.append(sel_label[i])
        label[i] = temp2
    
    return data_pad,file_out,sel_data,label

if __name__=="__main__":
    data_pad,file_out,sel_data,label=data_under_test(9000,9000)